
import os
import ntpath
import time
import argparse as ap
import requests
from tqdm import tqdm
from urllib.parse import quote
from loguru import logger
import concurrent.futures
from datetime import datetime
import re


class ImageDownloader:
    def __init__(self, targetname, startdate, enddate, focus=False,
                 ell_upper_limit=None, fwhm_upper_limit=None, workers=4,
                 api_version="V20250303",local=False,au_key=None):
        self.targetname = targetname
        self.local = local
        self.startdate = startdate
        self.enddate = enddate
        self.focus = focus
        self.ell_upper_limit = ell_upper_limit
        self.fwhm_upper_limit = fwhm_upper_limit
        self.workers = workers
        self.api_version = api_version
        self.base_url = f"http://192.168.16.70:8082/{self.api_version}/search_by_using_obj_filter"
        
        self.localurl="/all/data_products/"
        self.au_key=au_key

    def transform_urls(self,data):
        """转换URL中的端口到对应的存储节点路径"""
        transformed = data.copy()
        
        if self.api_version=="V20250303_KM":
            # 生成所有 27 个节点的映射规则
            port_mappings = {}
            for node_id in range(1, 30):
                # 生成端口号 (10001 到 10027)
                port = f"100{node_id:02d}"
                # 生成节点名称 (kmnode01 到 kmnode27)
                node_name = f"kmnode{node_id:02d}"
                # 正则：匹配类似 "http://任意IP:port"
                pattern = re.compile(rf'http://[^:/]+:{port}')
                # 替换后要拼接的路径
                replacement = f"/all/data_products/{node_name}/1m6"
                port_mappings[port] = (replacement, pattern)
            
            # 对每个字段逐个做替换，串联每一次的结果
            for key, value in transformed.items():
                if isinstance(value, str):
                    current = value  # 用一个局部变量保存「不断更新」的结果
                    for port, (replacement, pattern) in port_mappings.items():
                        # 注意这里改成用 current，而不是最初的 value
                        current = pattern.sub(replacement, current)
                    # 循环结束后，再把 current 赋回 transformed[key]
                    transformed[key] = current
            
            return transformed
        elif self.api_version.startswith("V20250303"):
    
            for key, value in transformed.items():
                if isinstance(value, str):
                    current = value

                    # 1. 替换统一的 16000 端口地址（固定 IP）
                    current = re.sub(
                        r"http://192\.168\.16\.206:16000/V20250303",
                        f"/home/mpilot/1m6/reception/V20250303",
                        current
                    )

                    # 2. 替换每个节点的 reception 数据（10001~10029）
                    for node_id in range(1, 30):
                        port = f"100{node_id:02d}"
                        pattern = re.compile(rf"http://[^:/]+:{port}/reception/V20250303")
                        replacement = f"/all/data_products/kmnode{node_id:02d}/1m6/reception/V20250303"
                        current = pattern.sub(replacement, current)

                    transformed[key] = current

            return transformed
        else:
            raise ValueError(f"Unsupported API version: {self.api_version}")
        
    def download_file(self, file_url, target_folder, file_type=None):
        os.makedirs(target_folder, exist_ok=True)
        filename = os.path.basename(file_url)

        if file_type == "fwhm_png":
            filename, ext = os.path.splitext(filename)
            filename = f"{filename}.fwhm{ext}"
        elif file_type == "preview_png":
            filename, ext = os.path.splitext(filename)
            filename = f"{filename}.preview{ext}"

        full_path = os.path.join(target_folder, filename)

        with requests.get(file_url, stream=True, timeout=20) as resp:
            if resp.status_code == 200:
                with open(full_path, 'wb') as f:
                    for chunk in resp.iter_content(chunk_size=8192):
                        f.write(chunk)
            else:
                # 记录错误日志
                error_msg = f"Download failed: {file_url} [HTTP {resp.status_code}]"
                logger.error(error_msg)
                # 抛出异常
                # 方法1: 抛出 HTTPError (需要传入 Response 对象)
                resp.raise_for_status()  # 这会根据状态码抛出对应的 HTTPError

        return full_path

    def query_and_download(self, target):
        localpath = os.path.join("images", target)
        os.makedirs(localpath, exist_ok=True)

        for band in ["g", "r","u", "v",  "i", "z"]:
            params = self._build_query_params(target, band)
            logger.info(f"Querying [{target}] in band [{band}]...")

            try:
                response = requests.get(self.base_url, params=params, timeout=20)
                response.raise_for_status()
                data = response.json()
                
            except Exception as e:
                logger.error(f"Failed to query {target} [{band}]: {e}")

                continue

            if data.get("code") != "200" or "data" not in data:
                logger.warning(f"No results for {target} [{band}]")
                continue

            log_file_path = os.path.join(localpath, f"images_{band}.list")
            with open(log_file_path, "w") as log_file:
                log_file.write("")

            for item in tqdm(data["data"], desc=f"{target}-{band}"):
                if item["status"] != "1":
                    logger.info(f"this item {item['sciimg']} is failed")
                    continue
                if self.local: #如果为真，则
                    
                    transformed = self.transform_urls(item)
                    logger.info(transformed)
                    os.system(f"stat {transformed['sciimg']}")
                    with open(log_file_path, "a") as log_file:
                        log_file.write(transformed['sciimg'] + "\n")
                        
                else:
                    try:
                        sciimg = item["sciimg"]
                        calmag = item["sciimg_gaia_sexfits"].replace("_gaia", "_CalMag")

                        if not self.focus and ("focus" in sciimg.lower() or "focus" in calmag.lower()):
                            continue

                        if not os.path.exists(os.path.join(localpath, calmag)):
                            path = self.download_file(calmag, localpath)
                            try:
                                from tools import utils
                                utils.CatalogueRegion(path)
                            except:
                                pass

                        if not os.path.exists(os.path.join(localpath, sciimg)):
                            self.download_file(sciimg, localpath)

                        with open(log_file_path, "a") as log_file:
                            log_file.write(ntpath.basename(sciimg) + "\n")

                    except Exception as e:
                        logger.error(f"[{target}] download error: {e}")
                        with open("failed_targets.txt", "a") as fail_log:
                            fail_log.write(f"{target}\n")

    def _build_query_params(self, target, band):
        encoded_target = quote(target)
        choose = "choose"

        if self.ell_upper_limit is None and self.fwhm_upper_limit is None:
            return {
                "object": encoded_target,
                "filter": band,
                "start_date": self.startdate,
                "end_date": self.enddate,
                "choose": choose,
                'au_key': self.au_key
            }

        # === Pixscale 说明 ===
        # 对于 i/z 波段：
        # - 2023-03-08 之前：使用 Andor CCD（像元 0.429 arcsec/pixel）
        # - 2023-03-08 ~ 2023-12-20：红通道没有相机（不应有有效 i/z 数据）
        # - 2023-12-21 起：启用 CCD290（像元 0.286 arcsec/pixel）

        pixscale = 0.429  # 默认使用 Andor 相机（早期）

        if band in ["i", "z"]:
            try:
                date_fmt = "%Y%m%d"
                dt_start = datetime.strptime(self.startdate, date_fmt)

                if dt_start >= datetime(2023, 12, 21):
                    pixscale = 0.286  # CCD290 新相机启用
                    logger.info(f"[{band}] {dt_start.date()} 使用 CCD290，相机像元=0.286")
                elif dt_start > datetime(2023, 3, 8):
                    logger.warning(f"[{band}] {dt_start.date()} 查询落在无红通道相机时间段（2023-03-08 至 2023-12-20），结果可能为空")
                else:
                    pixscale = 0.429  # 早期 Andor
                    logger.info(f"[{band}] {dt_start.date()} 使用 Andor CCD，pixscale=0.429")
            except Exception as e:
                logger.warning(f"日期解析失败，默认使用 pixscale=0.429: {e}")


        fwhm_pix = self.fwhm_upper_limit / pixscale
        return {
            "object": encoded_target,
            "filter": band,
            "start_date": self.startdate,
            "end_date": self.enddate,
            "medFWHM": f"{fwhm_pix:.2f}",
            "medEll": f"{self.ell_upper_limit:.2f}",
            "choose": choose,
            'au_key':self.au_key
        }

    def run(self):
        t_start = time.time()
        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 开始处理")

        if os.path.isfile(self.targetname):
            with open(self.targetname, 'r') as f:
                targets = [line.strip() for line in f if line.strip()]
            with concurrent.futures.ThreadPoolExecutor(max_workers=self.workers) as executor:
                futures = [executor.submit(self.query_and_download, t) for t in targets]
                for f in concurrent.futures.as_completed(futures):
                    f.result()
        else:
            self.query_and_download(self.targetname)

        elapsed = time.time() - t_start
        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] 完成，用时：{elapsed:.2f} 秒 ≈ {elapsed/60:.2f} 分钟")


def main():
    parser = ap.ArgumentParser(description="Image Downloader for Transients")
    parser.add_argument("-tn", "--targetname", type=str, default='sb24111301', help="目标名或文件路径")
    parser.add_argument("-sd", "--startdate", type=str, default='20241112', help="起始日期 yyyymmdd")
    parser.add_argument("-ed", "--enddate", type=str, default='20241115', help="终止日期 yyyymmdd")
    parser.add_argument("-f", "--focus", default=False, help="是否包含 focus 图像")
    parser.add_argument("-el", "--ell_upper_limit", type=float, default=None, help="椭率上限")
    parser.add_argument("-fw", "--fwhm_upper_limit", type=float, default=None, help="FWHM 上限（arcsec）")
    parser.add_argument("--workers", type=int, default=4, help="最大线程数")
    parser.add_argument("--api_version", type=str, default="V20250303_KM", help="接口版本路径（如 V20250303_KM）")
    parser.add_argument("--local", default=False, help="默认是通过网络下载，不读取本地磁盘")
    parser.add_argument("--au_key", default='2f0bb0f64a0b65ae0806d0fc295899f0', help="默认是通过网络下载，不读取本地磁盘")

    args = parser.parse_args()

    downloader = ImageDownloader(
        targetname=args.targetname,
        startdate=args.startdate,
        enddate=args.enddate,
        focus=args.focus,
        ell_upper_limit=args.ell_upper_limit,
        fwhm_upper_limit=args.fwhm_upper_limit,
        workers=args.workers,
        api_version=args.api_version,
        local=args.local,
        au_key=args.au_key
    )
    downloader.run()


if __name__ == '__main__':
    main()





# 下载地址:
# V20250303_KM: 昆明超算V20250303_KM版本的处理数据下载的地址 （推荐批量下载）
# V20250303: 丽江超算V20250303版本的处理数据下载的地址 （历史数据还未处理完成，不推荐批量下载）
# V20240815: 丽江超算V20240815版本的处理数据下载的地址 （不推荐）


"""
Image Downloader 脚本说明
=========================

用途：
    根据目标名或目标列表，从远程 API 下载天文图像文件（sciimg 和 CalMag）

基本用法：
    # 单目标（使用默认接口版本）
    python main_download_demo_img.py -tn grb250419a

    # 多目标文件（每行一个目标名），并发线程数设为6
    python main_download_demo_img.py -tn targets.txt --workers 6

    # 筛选图像质量（FWHM < 2.5 arcsec 且 椭率 < 0.3）
    python main_download_demo_img.py -tn targets.txt -fw 2.5 -el 0.3

    # 包含 focus 图像
    python main_download_demo_img.py -tn grb250419a --focus

    # 使用自定义 API 接口版本（如 V20250410_DEV）
    python main_download_demo_img.py -tn grb250419a --api_version V20250410_DEV

参数说明：
    -tn, --targetname          目标名或目标文件路径
    -sd, --startdate           查询起始日期（yyyymmdd）
    -ed, --enddate             查询结束日期（yyyymmdd）
    -f,  --focus               是否包含 focus 图像（默认不包含）
    -el, --ell_upper_limit     椭率上限
    -fw, --fwhm_upper_limit    FWHM 上限（单位：arcsec）
    --workers                  最大并发线程数（默认4）
    --api_version              查询接口路径（默认 V20250303_KM）

输出目录结构：
    images/
        └── <target>/
              ├── images_u.list
              ├── images_g.list
              ├── <sciimg>.fits
              ├── <CalMag>.fits
"""



# 字段介绍
# 字段名 | 含义（初步推测）
# exposure | 曝光时间（单位：秒）
# filename | 原始FITS文件下载链接
# filter | 滤光片波段（如 "v"）
# fwhm | 图像FWHM（Full Width Half Maximum，全宽半高，通常反映图像质量）
# fwhm_png | FWHM测量图像的PNG预览链接
# obs_date | 观测时间（UTC时间，格式："YYYY-MM-DD HH:MM:SS"）
# preview_png | 科学图像的整体预览PNG链接
# sciimg | 科学处理后的FITS图像（sciimg）下载链接
# sciimg_gaia_sexfits | 使用Gaia星表校准后的Source Extractor输出FITS文件链接
# sciimg_sexcat_gaia_XP | 使用Gaia星表和Gaia XP改进版输出的Source Extractor catalog链接
# sciimg_sexfits | Source Extractor直接输出的catalog FITS文件链接
# status | 任务状态（1通常表示成功，0可能表示失败或其他状态）
